# -*- coding: utf-8 -*-
"""
   File Name：     ocr.py
   Description :  wrapping for face recognition class, instantiation occurs in modules/face_server/utils/api.py
   Author :       KangJiaHui
   date：         2021/02/04
"""
import itertools
import re
import os
import yaml
from datetime import datetime, timedelta
from paddleocr import PaddleOCR
from module.utils.utils import base64_to_image, image_to_base64
from module.utils.error import AimPose, NameRecFail, NameMatchFail
from module.face_recog.faceRecognition import FaceRecognition
import cv2
import requests

face = FaceRecognition()


def resize_img(image, _max=1500):
    if image.shape[0] > _max or image.shape[1] > _max:
        if image.shape[0] > image.shape[1]:
            proportion = _max / image.shape[0]
        else:
            proportion = _max / image.shape[1]
        image = cv2.resize(image, None, fx=proportion, fy=proportion, interpolation=cv2.INTER_AREA)
    return image


def get_neighbor_box(lst, box, aim_pos, thresh):
    """
    From box location to find its neighbor, according to direction(FLAG) and distance(thresh).
    :param lst: list, return of get_all().
    :param box: list, e.x.[[1299.0, 480.0], [1448.0, 477.0], [1449.0, 514.0], [1300.0, 517.0]]
    :param aim_pos: str, in range of options: "RIGHT", "LEFT", "DOWN", "UP", "SAME"
    :param thresh: int, in mode "RIGHT", "LEFT", "DOWN" and "UP",
            thresh means how far the target could be away with keyword.
    :return: list, all lines matched. e.x.
    """
    result = [box]
    if aim_pos == "RIGHT":
        for line in lst:
            if 0 < abs(line[0][3][0] - box[0][2][0]) < thresh * 10 and 0 < abs(line[0][3][1] - box[0][2][1]) < thresh:
                result.append(line)
    elif aim_pos == "LEFT":
        for line in lst:
            if 0 < abs(line[0][2][0] - box[0][3][0]) < thresh * 10 and 0 < abs(line[0][2][1] - box[0][3][1]) < thresh:
                result.append(line)
    elif aim_pos == "DOWN":
        for line in lst:
            if 0 < abs(line[0][1][1] - box[0][2][1]) < thresh and \
                    0 < min(abs(line[0][2][0] - box[0][2][0]), abs(line[0][3][0] - box[0][3][0])) \
                    < abs(line[0][2][0] - line[0][3][0] - box[0][2][0] + box[0][3][0]) / 2:
                result.append(line)
    elif aim_pos == "UP":
        for line in lst:
            if 0 < abs(line[0][2][1] - box[0][1][1]) < thresh and \
                    0 < min(abs(line[0][2][0] - box[0][2][0]), abs(line[0][3][0] - box[0][3][0])) \
                    < abs(line[0][2][0] - line[0][3][0] - box[0][2][0] + box[0][3][0]) / 2:
                result.append(line)
    elif aim_pos == "SAME":
        pass
    else:
        raise AimPose
    return result


def get_location(keyword_regular, result):
    """
    Match the regular expression of keyword and give all locations for each matched keyword.
    :param result: list, return of get_all().
    :param keyword_regular: string, could be regular expression. e.x."姓名", "危险｜爆炸"
    :return: list, e.x.[[[[1299.0, 480.0], [1448.0, 477.0], [1449.0, 514.0], [1300.0, 517.0]], ('危险', 0.99250406)],
                    [[[1663.0, 480.0], [1886.0, 480.0], [1886.0, 510.0], [1663.0, 510.0]], ('爆炸', 0.8288886)]]
    """
    location = []
    for line in result:
        if re.search(keyword_regular, line[-1][0]):
            location.append(line)
    return location


def get_info_by_location(result, info_regular, keyword_location, aim_pos, thresh=30):
    # TODO 查找阈值像素范围需要根据图像像素大小自适应
    """
    Get information by keyword location.
    :param result: list, return of get_all().
    :param info_regular: the regular expression to match information.
    :param keyword_location: list, return from self.get_keywords_location()
            e.x.[[[1299.0, 480.0], [1448.0, 477.0], [1449.0, 514.0], [1300.0, 517.0]], ('危险', 0.99250406)]
    :param aim_pos: str, optional by following :["RIGHT", "LEFT", "DOWN", "UP", "SAME"]
            "RIGHT" means to find information in boxes at the right of keyword_location box.
            "LEFT", "DOWN", "UP" mean the same as "RIGHT".
            "SAME" means to find information in the same box as keyword_location.
    :param thresh: int, in mode "RIGHT", "LEFT", "DOWN" and "UP",
            thresh means how far the target could be away with keyword.
    :return: list, contains all matched info strings.
            e.x.['鲁FBR932', '鲁GLN851挂']
    """
    matched_info = []
    for box in keyword_location:
        target = get_neighbor_box(result, box, aim_pos, thresh)
        for t in target:
            info = re.search(info_regular, t[-1][0])
            if info:
                matched_info.append(info.group(0))
    return matched_info


def filter_one(result):
    """
    Do something with the first vision of result, such as date transform, invalid string delete.
    :param result: dict, the result dict for one paper.
        e.x. {"有效期": ["2018年4月", "2021年4月"], "经营范围": ["危险化学品"], "车牌": ["鲁FBR932"]}
    :return: dict, the processed result dict for one paper.
        e.x. {"有效期": "2021-4", "经营范围": "危险化学品", "车牌": "鲁FBR932"}
    """
    # print(result)
    for k, v in result.items():
        if k == "有效期" and v:
            result[k] = []
            for i in v:
                if len(i.split("-")) == 3:
                    date = i.split("-")
                    today = datetime.today().strftime("%Y-%m-%d").split("-")
                    if (str(date[0]) > str(today[0])) or \
                            (str(date[0]) == str(today[0]) and str(date[1]) > str(today[1])) or \
                            (str(date[0]) == str(today[0]) and str(date[1]) == str(today[1]) and
                             str(date[2]) > str(today[2])):
                        result[k].append(i)
                    # else:
                    #     result[k].append("已过期: " + i)
                if len(re.split("[年月日]", i)) == 3:
                    date = re.split("[年月]", i)[:2]
                    if 0 < int(date[1]) < 12:
                        date_str = date[0] + "-" + str(int(date[1]) + 1) + "-01"
                    elif int(date[1]) == 12:
                        date_str = str(int(date[0]) + 1) + "-" + str(1) + "-01"
                    else:
                        continue
                    validity = datetime.strptime(date_str, '%Y-%m-%d') + timedelta(days=-1)
                    date_str_complete = validity.strftime("%Y-%m-%d")
                    today = datetime.today().strftime("%Y-%m-%d").split("-")
                    if (str(date[0]) > str(today[0])) or \
                            (str(date[0]) == str(today[0]) and str(date[1]) >= str(today[1])):
                        result[k].append(date_str_complete)
                    # else:
                    #     result[k].append("已过期: " + date_str_complete)
                elif len(re.split("[年月日]", i)) == 4:
                    date = re.split("[年月日]", i)[:3]
                    today = datetime.today().strftime("%Y-%m-%d").split("-")
                    if (str(date[0]) > str(today[0])) or \
                            (str(date[0]) == str(today[0]) and str(date[1]) > str(today[1])) or \
                            (str(date[0]) == str(today[0]) and str(date[1]) == str(today[1]) and
                             str(date[2]) > str(today[2])):
                        result[k].append(date[0] + "-" + date[1] + "-" + date[2])
                    # else:
                    #     result[k].append("已过期: " + date[0] + "-" + date[1] + "-" + date[2])
            v = result[k]
        if k == "姓名" and v:
            v = result[k] = [i for i in v if not re.search("姓名|性别|中国|国籍", i)]
        if len(v) == 0:
            result[k] = None
        if type(result[k]) == list:
            v = list(set(v))
        # if len(v) == 1:
            result[k] = v[0]

    return result


class OCR(object):
    def __init__(self):
        self.model = PaddleOCR(use_angle_cls=True, lang="ch")
        # self.my_model = PaddleOCR(det_model_dir='module/params/my_ch_db_res18',
        #                           rec_model_dir='module/params/ch_ppocr_mobile_v2.0_rec_infer',
        #                           rec_char_dict_path='module/params/ppocr_keys_v1.txt',
        #                           cls_model_dir='module/params/ch_ppocr_mobile_v2.0_cls_infer',
        #                           use_angle_cls=True)
        with open(os.path.join(os.getcwd(), "confs/certificate.yaml"), "r", encoding='UTF-8') as f:
            self.config = yaml.load(f, Loader=yaml.FullLoader)

    def get_all(self, img):
        """
        Recognize all information from an image.
        :param img: numpy.ndarray
        :return: list, e.x.[[[[1319.0, 171.0], [1673.0, 173.0], [1673.0, 232.0], [1319.0, 230.0]], ('服务单位', 0.9971724)]
        , [[[333.0, 201.0], [414.0, 201.0], [414.0, 248.0], [333.0, 248.0]], ('姓名', 0.9986272)]]
        """
        result = self.model.ocr(img, cls=True)
        # print(result)
        return result

    def name_recog(self, result):
        """
        Recognize witch certificate it is, according to the OCR result of image.
        :param result: list, the OCR result of input image.
        :return: string, the key of matched certificate.
        """
        for key in self.config:
            regular = self.config[key]["证件名"]
            if isinstance(regular, list):
                if list({'RIGHT', 'LEFT', 'DOWN', 'UP', 'SAME'} & set(regular)):
                    location = get_location(regular[0], result)
                    if get_info_by_location(result, regular[2], location, regular[1], 100):
                        return key
                else:
                    flag = True
                    for i in regular:
                        # print(get_location(i, result))
                        if not get_location(i, result):
                            flag = False
                    if flag:
                        return key
            elif isinstance(regular, str):
                if get_location(regular, result):
                    return key
        return None

    def recog_one(self, paper_name, image_base64, face_flag):
        """
        Recognize one image, and to update tmp_result dict
        :param paper_name: str, the name of certificate. e.x. "道路危险货物运输驾驶员证"
        :param face_flag: bool, if True, extract face feature from input image.
        :param image_base64: image encoded in base64
        :return: result: e.x. {"证件名": "中华人民共和国道路运输证", "有效期": None, "经营范围": "危险化学品", "车牌": "鲁FBR932"}
        """
        if re.search("行驶证|驾驶证", paper_name):
        # if False:
            image = base64_to_image(image_base64)
            image = resize_img(image)
            image_base64 = image_to_base64(image)
            result = self.baidu_ocr(paper_name, image_base64)
        else:
            image = base64_to_image(image_base64)
            image = resize_img(image)
            ocr_result = self.get_all(image)
            # print(ocr_result)
            name = self.name_recog(ocr_result)
            if not name:
                raise NameRecFail
            if name != paper_name:
                raise NameMatchFail
            result = {"证件名": name}
            for key in self.config[name]:
                if key == "证件名":
                    continue
                method = self.config[name][key]
                if isinstance(method, list):
                    loc = get_location(method[0], ocr_result)
                    result[key] = get_info_by_location(ocr_result, method[2], loc, method[1], 30)
                elif isinstance(method, str):
                    res = []
                    for line in ocr_result:
                        info = re.search(method, line[-1][0])
                        if info:
                            res.append(info.group(0))
                    # print(res)
                    result[key] = res
            filter_one(result)
        if face_flag:
            feature, bbox, face_img = face.face_register(image_base64, 0.5)
            result["人脸"] = feature
            result["人脸位置"] = bbox  # [left, top, right, bottom]
            result["切图"] = face_img
        return result

    def face_match(self, image_base64, face_feature):
        """
        Match face in input image with the faces in all certificates. Once a certificate matched, return True.
        :param face_feature:
        :param image_base64: image encoded in base64
        :return: bool, True means face matched, False means face not matched.
        """
        return face.match_identity(face_feature, 0.5, 0.5, image_base64)

    def car_plate_extract(self, image_base64):
        """
        Extract all car plates in input image.
        :param image_base64: image encoded in base64
        :return: str, the express of detected plate.
        """
        image = base64_to_image(image_base64)
        image = resize_img(image)
        ocr_result = self.get_all(image)
        plate = []
        for line in ocr_result:
            info = re.search('(^[\u4E00-\u9FA5][A-Z][A-Z0-9]{4}[A-Z0-9挂学警港澳]$)|([A-Z0-9]{4}[挂])', line[-1][0])
            if info:
                bbox = line[0]
                plate.append([info.group(0), bbox])
        return plate

    @staticmethod
    def face_register(image_base64):
        """
        Extract face feature vector of input image.
        :param image_base64: image encoded in base64
        :return: list, face feature vector of input image.
        """
        return face.face_register(image_base64, 0.5)

    @staticmethod
    def baidu_ocr(paper_name, image_base64):
        params = {"image": image_base64}
        access_token = '24.462c64a56f834db10ca99588b6dab524.2592000.1625977451.282335-24121783'
        headers = {'content-type': 'application/x-www-form-urlencoded'}
        if re.search("驾驶证", paper_name):
            request_url = "https://aip.baidubce.com/rest/2.0/ocr/v1/driving_license"
            request_url = request_url + "?access_token=" + access_token
            response = requests.post(request_url, data=params, headers=headers).json()
            result = {"证件名": paper_name, "姓名": response['words_result']['姓名']['words'],
                      "证号": response['words_result']['证号']['words'],
                      "准驾车型": response['words_result']['准驾车型']['words']}
            if "至" in response['words_result'].keys():
                result["有效期"] = response['words_result']['至']['words'][:4] + '-' + \
                                response['words_result']['至']['words'][4:6] + '-' + \
                                response['words_result']['至']['words'][6:]
            elif "有效起始日期" in response['words_result'].keys() and "有效期限" in response['words_result'].keys():
                year = int(response['words_result']['有效起始日期']['words'][:4]) + \
                       int(re.search("[0-9]+", response['words_result']['有效期限']['words']).group(0))
                result["有效期"] = str(year) + '-' + response['words_result']['有效起始日期']['words'][4:6] + '-' + \
                                response['words_result']['有效起始日期']['words'][6:]
        elif re.search("行驶证", paper_name):
            request_url = "https://aip.baidubce.com/rest/2.0/ocr/v1/vehicle_license"
            request_url = request_url + "?access_token=" + access_token
            response = requests.post(request_url, data=params, headers=headers).json()
            result = {"证件名": paper_name, "有效期": None, "车牌": response['words_result']['号牌号码']['words'],
                      "使用性质": response['words_result']['使用性质']['words']}
        else:
            raise NameMatchFail
        return result
